def createDataSet():
    dataSet=[[1,1,'yes'],
             [1,1,'yes'],
             [1,0,'no'],
             [0,1,'no'],
             [0,1,'no']]
    labels=['no surfacing','flippers']
    return dataSet,labels

from math import log
def calcShannonEnt(dataSet):
    numEntries=len(dataSet)
    lablesCounts={}
    for featVec in dataSet:
        currentLable=featVec[-1]
        if currentLable not in lablesCounts.keys():
            lablesCounts[currentLable]=0
        lablesCounts[currentLable]+=1
    shannonEnt=0.0
    for key in lablesCounts:
        prob=float(lablesCounts[key])/numEntries
        shannonEnt-=prob*log(prob,2)
    return shannonEnt

def splitDataSet(dataSet,axis,value):
    retDataSet=[]
    for featVec in dataSet:
        if featVec[axis]==value:
            reduceFeatVec=featVec[:axis]
            reduceFeatVec.extend(featVec[axis+1:])
            retDataSet.append(reduceFeatVec)
    return retDataSet

def chooseBestFeatureToSpilt(dataSet):
    numFeatures=len(dataSet[0])-1
    baseEntropy=calcShannonEnt(dataSet)
    bestInfoGain=0.0;bestFeature=-1;
    for i in range(numFeatures):
        feaList=[example[i] for example in dataSet]
        uniqueValues=set(feaList)
        newEntropy=0.0
        for value in uniqueValues:
            subDataSet=splitDataSet(dataSet,i,value)
            prob=len(subDataSet)/float(len(dataSet))
            newEntropy+=prob*calcShannonEnt(subDataSet)
        infoGain=baseEntropy-newEntropy
        if infoGain>bestInfoGain:
            bestInfoGain=infoGain
            bestFeature=i
    return bestFeature

import operator

def majorityCnt(classList):
    classCnt={}
    for vote in classList:
        if vote not in classCnt.keys():
            classCnt[vote]=0
        classCnt[vote]+=1
    sortedClassCnt=sorted(classCnt.items(),key=operator.itemgetter(1),reverse=True)
    return sortedClassCnt[0][0]

def createTree(dataSet,labels):
    classList=[example[-1] for example in dataSet]
    classCnt=classList.count(classList[0])
    if classCnt==len(classList):
        return classList[0]
    if len(dataSet[0])==1:
        return majorityCnt(classList)
    bestFeat=chooseBestFeatureToSpilt(dataSet)
    bestFeatLable=labels[bestFeat]
    myTree={bestFeatLable:{}}
    del(labels[bestFeat])
    featValues=[example[bestFeat] for example in dataSet]
    uniqueVals=set(featValues)
    for value in uniqueVals:
        subLabels=labels[:]
        myTree[bestFeatLable][value]=createTree(splitDataSet(dataSet,bestFeat,value),subLabels)
    return myTree

import matplotlib.pyplot as plt

# 定义决策节点和叶节点的样式
decisionNode = dict(boxstyle="sawtooth", fc="0.8")
leafNode = dict(boxstyle="round4", fc="0.8")
arrow_args = dict(arrowstyle="<-")

# 计算叶节点数
def getNumLeafs(myTree):
    numLeafs = 0
    firstStr = list(myTree.keys())[0]
    secondDict = myTree[firstStr]
    for key in secondDict.keys():
        if type(secondDict[key]) == dict:  # 如果是字典类型，说明是内部节点
            numLeafs += getNumLeafs(secondDict[key])
        else:
            numLeafs += 1  # 叶节点计数
    return numLeafs

# 计算树的深度
def getTreeDepth(myTree):
    maxDepth = 0
    firstStr = list(myTree.keys())[0]
    secondDict = myTree[firstStr]
    for key in secondDict.keys():
        if type(secondDict[key]) == dict:
            thisDepth = 1 + getTreeDepth(secondDict[key])
        else:
            thisDepth = 1
        if thisDepth > maxDepth:
            maxDepth = thisDepth
    return maxDepth

# 绘制节点
def plotNode(nodeTxt, centerPt, parentPt, nodeType):
    createPlot.ax1.annotate(nodeTxt, xy=parentPt, xycoords='axes fraction',
                            xytext=centerPt, textcoords='axes fraction',
                            va="center", ha="center", bbox=nodeType,
                            arrowprops=arrow_args)

# 在父子节点间添加文本
def plotMidText(entrPt, parentPt, txtString):
    xMid = (parentPt[0] + entrPt[0])/2.0
    yMid = (parentPt[1] + entrPt[1])/2.0
    createPlot.ax1.text(xMid, yMid, txtString, va="center", ha="center")

# 递归绘制树
def plotTree(myTree, parentPt, nodeTxt):
    numLeafs = getNumLeafs(myTree)
    depth = getTreeDepth(myTree)
    firstStr = list(myTree.keys())[0]
    
    # 计算当前节点位置
    cntrPt = (plotTree.xOff + (1.0 + numLeafs)/2.0/plotTree.totalW, plotTree.yOff)
    
    # 绘制连接线和文本
    plotMidText(cntrPt, parentPt, nodeTxt)
    plotNode(firstStr, cntrPt, parentPt, decisionNode)
    
    secondDict = myTree[firstStr]
    plotTree.yOff -= 1.0/plotTree.totalD  # 下移层级
    
    for key in secondDict:
        if isinstance(secondDict[key], dict):  # 判断是否是子树
            plotTree(secondDict[key], cntrPt, str(key))
        else:  # 叶节点处理
            plotTree.xOff += 1.0/plotTree.totalW
            plotNode(secondDict[key], (plotTree.xOff, plotTree.yOff), cntrPt, leafNode)
            plotMidText((plotTree.xOff, plotTree.yOff), cntrPt, str(key))
    
    plotTree.yOff += 1.0/plotTree.totalD  # 回溯时恢复层级

# 主绘制函数
def createPlot(inTree):
    fig = plt.figure(1, facecolor='white', figsize=(8, 4))
    fig.clf()
    axprops = dict(xticks=[], yticks=[])
    createPlot.ax1 = plt.subplot(111, frameon=False, **axprops)
    
    # 初始化全局变量
    plotTree.totalW = float(getNumLeafs(inTree))
    plotTree.totalD = float(getTreeDepth(inTree))
    plotTree.xOff = -0.5/plotTree.totalW
    plotTree.yOff = 1.0
    
    plotTree(inTree, (0.5, 1.0), '')
    plt.show()

def classify(inputTree, featLabels, testVec):
    firstStr = list(inputTree.keys())[0]          # 修正 keys() 访问方式
    secondDict = inputTree[firstStr]              # 修正字典访问为方括号
    featIndex = featLabels.index(firstStr)
    
    for key in secondDict.keys():
        if testVec[featIndex] == key:
            if type(secondDict[key]).__name__ == 'dict':  # 修正类型判断语法
                classLabel = classify(secondDict[key], featLabels, testVec)
            else:
                classLabel = secondDict[key]      # 修正缩进和冒号
            return classLabel    
        
def storeTree(inputTree, filename):
    import pickle
    import json, gzip
 #   with gzip.open(filename, 'wb') as f:
 #   with open(filename, 'wb') as f:  # 修正：使用二进制写入模式，统一变量名，改用with语句
 #       pickle.dump(inputTree, f)
 #       json.dump(inputTree, f)
    with open("data.json", "w", encoding="utf-8") as f:
        json.dump(inputTree, f, indent=4)  # indent 参数美化格式

def grabTree(filename):
    import pickle
    import json, gzip
#    with gzip.open(filename, 'rb') as f:
#    with open(filename, 'rb') as f:  # 修正：使用二进制读取模式，改用with语句
#        return pickle.load(f)
#         return json.load(f)
# 读取字典
    with gzip.open("data.json", "r", encoding="utf-8") as f:
        loaded_data = json.load(f)
    return loaded_data

    

if __name__ == "__main__":
    myDat,labels=createDataSet()
    print(myDat)
    print(labels)
    newLabels=labels.copy()
    print(calcShannonEnt(myDat))
    print(splitDataSet(myDat,0,1))
    print(splitDataSet(myDat,0,0))
    print(splitDataSet(myDat,1,0))
    print(chooseBestFeatureToSpilt(myDat))
    mytree=createTree(myDat,labels)
    print(mytree)
    print(getNumLeafs(mytree))
    print(getTreeDepth(mytree))
    createPlot(mytree)
    storeTree(mytree,"dt.txt")
    mytree=grabTree("dt.txt")
    print(classify(mytree,newLabels,[1,0]))
